/*
 * Copyright (C) 2008 by The Regents of the University of California
 * Redistribution of this file is permitted under the terms of the GNU
 * Public License (GPL).
 *
 * @author Junghoo "John" Cho <cho AT cs.ucla.edu>
 * @date 3/24/2008
 */
 
#include "BTreeIndex.h"
#include "BTreeNode.h"

using namespace std;

struct IndexHeaderContents
{
    PageId rootPid;
    int treeHeight;
};

union IdxHeader
{
    struct IndexHeaderContents meta;
    char buffer[PageFile::PAGE_SIZE];
};

IdxHeader gl_idxHeaderTemplate;

/*
 * BTreeIndex constructor
 */
BTreeIndex::BTreeIndex()
{
    rootPid = -1;
    treeHeight = -1;
    last_leafPid = -1;
    status = MODE_CLOSED;
}

/*
 * Open the index file in read or write mode.
 * Under 'w' mode, the index file should be created if it does not exist.
 * @param indexname[IN] the name of the index file
 * @param mode[IN] 'r' for read, 'w' for write
 * @return error code. 0 if no error
 */
RC BTreeIndex::open(const string& indexname, char mode)
{
    if(status != MODE_CLOSED)
    {
        // @ERROR: Cannot open an already open index file
        
        return(RC_FILE_OPEN_FAILED);
    }
    
    switch(mode)
    {
        case('r'):
        {
            // Open the index file for reading
            
            if(pf.open(indexname, 'r') != 0)
            {
                // @ERROR: Could not open pagefile
                
                status = MODE_CLOSED;
                return(RC_FILE_OPEN_FAILED);
            }
                
            status = MODE_READ;
            break;
        }
        case('w'):
        {
            // Open the index file for writing
            
            if(pf.open(indexname, 'w') != 0)
            {
                // @ERROR: Could not open pagefile
                
                status = MODE_CLOSED;
                return(RC_FILE_OPEN_FAILED);
            }
            
            status = MODE_WRITE;
            break;
        }
        default:
        {
            // @ERROR: Invalid mode
            
            status = MODE_CLOSED;
            return(RC_FILE_OPEN_FAILED);
        }
    }
    
    char buffer[PageFile::PAGE_SIZE];
    
    if(pf.endPid() == 0)
    {
        // A new index file was created
        
        memset(buffer, 0, PageFile::PAGE_SIZE / sizeof(char));
        memset(buffer, -1, sizeof(PageId));
        memset(&buffer[4], -1, sizeof(treeHeight));
        
        if(pf.write(0, buffer) != 0)
        {
            // @ERROR: Could not create index header page
            
            status = MODE_CLOSED;
            return(RC_FILE_OPEN_FAILED);
        }
        
        rootPid = -1;
        treeHeight = 0;
    }
    else
    {
        // Retrieve index file information from header page
        
        if(pf.read(0, buffer) != 0)
        {
            // @ERROR: Could not retrieve index header page
            
            status = MODE_CLOSED;
            return(RC_FILE_OPEN_FAILED);
        }
        
        memcpy(&rootPid, buffer, sizeof(rootPid));
        memcpy(&treeHeight, &buffer[4], sizeof(treeHeight));
    }
    
    return(0);
}

/*
 * Close the index file.
 * @return error code. 0 if no error
 */
RC BTreeIndex::close()
{
    if(status == MODE_CLOSED)
    {
        // @ERROR: Cannot close an already closed index file
        
        return(RC_FILE_CLOSE_FAILED);
    }
    else if(status == MODE_READ)
    {
        // Closing a read-only index file
        
        rootPid = -1;
        treeHeight = -1;
        
        return(0);
    }
    
    // Flush the index buffer to the file
    char buffer[PageFile::PAGE_SIZE];
    
    memcpy(buffer, &rootPid, sizeof(PageId));
    memcpy(&buffer[4], &treeHeight, sizeof(treeHeight));
    
    if(pf.write(0, buffer) != 0)
    {
        // @ERROR: Could not create index header page
        
        return(RC_FILE_CLOSE_FAILED);
    }
    
    rootPid = -1;
    treeHeight = -1;
    
    return(0);
}

/*
 * Insert (key, RecordId) pair to the index.
 * @param key[IN] the key for the value inserted into the index
 * @param rid[IN] the RecordId for the record being inserted into the index
 * @return error code. 0 if no error
 */
RC BTreeIndex::insert(int key, const RecordId& rid)
{
    if(status != MODE_WRITE)
    {
        // @ERROR: The index file does not have write permissions
        
        return(RC_MODE_INVALID);
    }
    
    BTLeafNode leaf;
    BTNonLeafNode nleaf;
    
    if(rootPid < 1)
    {
        // The index file contains no nodes
        
        // Set the root page ID to reference the next available page ID
        rootPid = pf.endPid();
        
        // Insert the key into the newly created leaf node
        leaf.insert(key, rid);
        leaf.write(rootPid, pf);
        
        treeHeight = 1;
        
        return(0);
    }
    
    int of_key;
    int of_pid;
    
    PageId sibPid;
    
    switch(insertIndex(key, rid, rootPid, 1, of_key, of_pid))
    {
        case(0):
        {
            break;
        }
        case(RC_LEAF_OVERFLOW):
        {
            BTLeafNode leaf;
            BTLeafNode sib_leaf;
            
            int sib_key;
            PageId sib_pid;
            int sib_eid;
            
            if(leaf.read(rootPid, pf) != 0)
            {
                // @ERROR: Could not retrieve leaf node from pagefile
                
                return(RC_INSERT_FAILED);
            }
            
            if(leaf.insertAndSplit(key, rid, sib_leaf, sib_key) != 0)
            {
                // @ERROR: Could not insert key into leaf node
                
                return(RC_INSERT_FAILED);
            }
            
            // The key was successfully inserted into a leaf node
            
            sibPid = pf.endPid();
            
            if(leaf.write(rootPid, pf) != 0)
            {
                // @ERROR: Could not write leaf node contents to 
                //         pagefile
                
                return(RC_INSERT_FAILED);
            }
            
            if(sib_leaf.write(sibPid, pf) != 0)
            {
                // @ERROR: Could not write sibling leaf node 
                //         contents to pagefile
                
                return(RC_INSERT_FAILED);
            }
            
            if(nleaf.initializeRoot(rootPid, sib_key, sibPid) != 0)
            {
                // The first key in the sibling leaf node could not 
                // be added to the nonleaf node because the nonleaf
                // node was full
                
                return(RC_INSERT_FAILED);
            }
            
            // The first key in the sibling leaf node was 
            // successfully inserted into the nonleaf node
            
            rootPid = pf.endPid();
            
            if(nleaf.write(rootPid, pf) != 0)
            {
                // @ERROR: Could not write nonleaf node contents to 
                //         pagefile
                
                return(RC_INSERT_FAILED);
            }
            
            treeHeight++;
            
            gl_idxHeaderTemplate.meta.rootPid = rootPid;
            gl_idxHeaderTemplate.meta.treeHeight = treeHeight;
            
            pf.write(0, gl_idxHeaderTemplate.buffer);
            
            return(0);
        }
        case(RC_NLEAF_OVERFLOW):
        {
            BTNonLeafNode child_nleaf;
            BTNonLeafNode childsib_nleaf;
            
            int mid_key;
            
            if(child_nleaf.read(rootPid, pf) != 0)
            {
                // @ERROR: Could not retrieve nonleaf node from pagefile
                
                return(RC_INSERT_FAILED);
            }
            
            if(child_nleaf.insertAndSplit(of_key, of_pid, childsib_nleaf, 
                                         mid_key) != 0)
            {
                // @ERROR: Could not insert key into nonleaf node
                
                return(RC_INSERT_FAILED);
            }
            
            if(child_nleaf.write(rootPid, pf) != 0)
            {
                // @ERROR: Could not write nonleaf node contents to 
                //         pagefile
                
                return(RC_INSERT_FAILED);
            }
            
            sibPid = pf.endPid();
            
            if(childsib_nleaf.write(sibPid, pf) != 0)
            {
                // @ERROR: Could not write nonleaf node contents to 
                //         pagefile
                
                return(RC_INSERT_FAILED);
            }
            
            if(nleaf.initializeRoot(rootPid, mid_key, sibPid) != 0)
            {
                // Could not insert the new root key entry into the new root 
                // nonleaf node
                
                return(RC_INSERT_FAILED);
            }
            
            // The first key in the sibling nonleaf node was 
            // successfully inserted into the nonleaf node
            
            rootPid = pf.endPid();
            
            if(nleaf.write(rootPid, pf) != 0)
            {
                // @ERROR: Could not write nonleaf node contents to 
                //         pagefile
                
                return(RC_INSERT_FAILED);
            }
            
            treeHeight++;
            
            gl_idxHeaderTemplate.meta.rootPid = rootPid;
            gl_idxHeaderTemplate.meta.treeHeight = treeHeight;
            
            pf.write(0, gl_idxHeaderTemplate.buffer);
            
            break;
        }
        case(RC_INSERT_FAILED):
        default:
        {
            // @ERROR: Could not insert key into index
            
            return(RC_INSERT_FAILED);
        }
    }
    
    return(0);
}

/*
 * Find the leaf-node index entry whose key value is larger than or 
 * equal to searchKey, and output the location of the entry in IndexCursor.
 * IndexCursor is a "pointer" to a B+tree leaf-node entry consisting of
 * the PageId of the node and the SlotID of the index entry.
 * Note that, for range queries, we need to scan the B+tree leaf nodes.
 * For example, if the query is "key > 1000", we should scan the leaf
 * nodes starting with the key value 1000. For this reason,
 * it is better to return the location of the leaf node entry 
 * for a given searchKey, instead of returning the RecordId
 * associated with the searchKey directly.
 * Once the location of the index entry is identified and returned 
 * from this function, you should call readForward() to retrieve the
 * actual (key, rid) pair from the index.
 * @param key[IN] the key to find.
 * @param cursor[OUT] the cursor pointing to the first index entry
 *                    with the key value.
 * @return error code. 0 if no error.
 */
RC BTreeIndex::locate(int searchKey, IndexCursor& cursor)
{
    if(status == MODE_CLOSED)
    {
        // @ERROR: The index file is not open
        
        return(RC_FILE_OPEN_FAILED);
    }
    
    if((rootPid < 1) || (treeHeight < 0))
    {
        // @ERROR: The index file data structure is invalid
        
        return(RC_DATA_FAULT);
    }
    
    int heightCursor = 1;
    PageId pid = rootPid;
    
    while(heightCursor < treeHeight)
    {
        // Currently processing at a nonleaf level
        
        BTNonLeafNode nleaf;
        int eid;
        
        if(nleaf.read(pid, pf) != 0)
        {
            // @ERROR: Could not read nonleaf node pagefile
            
            return(RC_FILE_READ_FAILED);
        }
        
        if(nleaf.locateChildPtr(searchKey, pid, eid) != 0)
        {
            // @ERROR: Could not locate child pointer associated with the 
            //         search key
            
            return(RC_KEY_NOT_FOUND);
        }
        
        heightCursor++;
    }
    
    // Currently processing at the leaf level
    
    BTLeafNode leaf;
    int keyCount;
    int eid;
    
    if(leaf.read(pid, pf) != 0)
    {
        // @ERROR: Could not read leaf node pagefile
        
        return(RC_FILE_READ_FAILED);
    }
    
    if(leaf.locate(searchKey, eid) != 0)
    {
        // @ERROR: Could not locate search key in leaf node entry
        
        return(RC_KEY_NOT_FOUND);
    }
    
    cursor.pid = pid;
    cursor.eid = eid;
    
    last_leaf = leaf;
    last_leafPid = pid;
    
    return(0);
}

/*
 * Read the (key, rid) pair at the location specified by the index cursor,
 * and move foward the cursor to the next entry.
 * @param cursor[IN/OUT] the cursor pointing to an leaf-node index entry in the b+tree
 * @param key[OUT] the key stored at the index cursor location.
 * @param rid[OUT] the RecordId stored at the index cursor location.
 * @return error code. 0 if no error
 */
RC BTreeIndex::readForward(IndexCursor& cursor, int& key, RecordId& rid)
{
    if((cursor.pid < 0) || (cursor.pid >= pf.endPid()))
    {
        // @ERROR: The index cursor page ID is invalid
        
        return(RC_INVALID_PAGEID);
    }
    else if((cursor.eid < 0) || (cursor.eid >= BTLeafNode::MAX_NODE_ENTRIES))
    {
        // @ERROR: The index cursor entry ID is invalid
        
        return(RC_INVALID_ENTRYID);
    }
    else if(cursor.pid != last_leafPid)
    {
        // The cached leaf node does not reference the correct page
        
        if(last_leaf.read(cursor.pid, pf) != 0)
        {
            // @ERROR: Could not read associated page to the leaf node
            
            return(RC_FILE_READ_FAILED);
        }
    }
    
    // Reads the leaf node entry at the entry ID specified
    last_leaf.readEntry(cursor.eid, key, rid);
    
    cursor.eid++;
    
    if(cursor.eid >= last_leaf.getKeyCount())
    {
        // All the entries in the current leaf node have been processed
        
        // Get the page ID for the next leaf node
        cursor.pid = last_leaf.getNextNodePtr();
        cursor.eid = 0;
        
        // Update the leaf node cache to reflect the new leaf node
        if(last_leaf.read(cursor.pid, pf) != 0)
        {
            // @ERROR: Could not read associated page to the leaf node
            
            return(RC_FILE_READ_FAILED);
        }
        
        last_leafPid = cursor.pid;
    }
    
    return(0);
}

// @USR: User-defined methods...
RC BTreeIndex::insertIndex(int key, const RecordId &rid, PageId pid, 
                           int height, int &of_key, PageId &of_pid)
{
    if((pid <= 0) || (pid >= pf.endPid()))
    {
        // @ERROR: The index cursor page ID is invalid
        
        return(RC_INVALID_PAGEID);
    }
    else if(height > treeHeight)
    {
        // @ERROR: Invalid height parameter
        
        return(RC_INVALID_ARG);
    }
    
    if(height == treeHeight)
    {
        // Current processing at the leaf node level
        
        BTLeafNode leaf;
        
        if(leaf.read(pid, pf) != 0)
        {
            // @ERROR: Could not retrieve leaf node from pagefile
            
            return(RC_INSERT_FAILED);
        }
        
        if(leaf.insert(key, rid) == 0)
        {
            // The key was successfully inserted into the leaf node
            
            if(leaf.write(pid, pf) != 0)
            {
                // @ERROR: Could not write leaf node contents to pagefile
                
                return(RC_INSERT_FAILED);
            }
            
            return(0);
        }
        
        // The key could not be added to the leaf node because the leaf node 
        // was full
        
        return(RC_LEAF_OVERFLOW);
    }
    else
    {
        // Current processing at a nonleaf node level
        
        BTNonLeafNode nleaf;
        
        if(nleaf.read(pid, pf) != 0)
        {
            // @ERROR: Could not retrieve nonleaf node from pagefile
            
            return(RC_INSERT_FAILED);
        }
        
        PageId childPid;
        PageId sibPid;
        int eid;
        
        if(nleaf.locateChildPtr(key, childPid, eid) != 0)
        {
            // @ERROR: Current nonleaf node is invalid and contains no entries
            
            return(RC_INSERT_FAILED);
        }
        
        // Traverse through the next level of the index tree structure
        switch(insertIndex(key, rid, childPid, height + 1, of_key, of_pid))
        {
            case(0):
            {
                // Key was successfully added to the index
                
                return(0);
            }
            case(RC_LEAF_OVERFLOW):
            {
                // The previous operation encountered a leaf overflow
                
                BTLeafNode leaf;
                BTLeafNode sib_leaf;
                
                int sib_key;
                PageId sib_pid;
                int sib_eid;
                
                if(leaf.read(childPid, pf) != 0)
                {
                    // @ERROR: Could not retrieve leaf node from pagefile
                    
                    return(RC_INSERT_FAILED);
                }
                
                if(leaf.insertAndSplit(key, rid, sib_leaf, sib_key) != 0)
                {
                    // @ERROR: Could not insert key into leaf node
                    
                    return(RC_INSERT_FAILED);
                }
                
                // The key was successfully inserted into a leaf node
                
                sibPid = pf.endPid();
                
                if(leaf.write(childPid, pf) != 0)
                {
                    // @ERROR: Could not write leaf node contents to 
                    //         pagefile
                    
                    return(RC_INSERT_FAILED);
                }
                
                if(sib_leaf.write(sibPid, pf) != 0)
                {
                    // @ERROR: Could not write sibling leaf node 
                    //         contents to pagefile
                    
                    return(RC_INSERT_FAILED);
                }
                
                if(nleaf.nleaf_pf.node.entryCount == 
                            BTNonLeafNode::MAX_NODE_ENTRIES)
                {
                    // The nonleaf node is full
                    
                    of_key = sib_key;
                    of_pid = sibPid;
                    
                    return(RC_NLEAF_OVERFLOW);
                }
                else
                {
                    // The key can be inserted into the nonleaf node
                    
                    if(nleaf.insert(sib_key, sibPid) != 0)
                    {
                        // @ERROR: Could not insert key into nonleaf node
                        
                        return(RC_INSERT_FAILED);
                    }
                }
                
                // The first key in the sibling leaf node was 
                // successfully inserted into the nonleaf node
                
                if(nleaf.write(pid, pf) != 0)
                {
                    // @ERROR: Could not write nonleaf node contents to 
                    //         pagefile
                    
                    return(RC_INSERT_FAILED);
                }
                
                return(0);
            }
            case(RC_NLEAF_OVERFLOW):
            {
                // The previous operation encountered a nonleaf overflow
                
                BTNonLeafNode child_nleaf;
                BTNonLeafNode childsib_nleaf;
                int mid_key;
                
                if(child_nleaf.read(childPid, pf) != 0)
                {
                    // @ERROR: Could not retrieve nonleaf node from pagefile
                    
                    return(RC_INSERT_FAILED);
                }
                
                if(child_nleaf.insertAndSplit(of_key, of_pid, childsib_nleaf, 
                                             mid_key) != 0)
                {
                    // @ERROR: Could not insert key into nonleaf node
                    
                    return(RC_INSERT_FAILED);
                }
                
                sibPid = pf.endPid();
                
                if(child_nleaf.write(childPid, pf) != 0)
                {
                    // @ERROR: Could not write nonleaf node contents to 
                    //         pagefile
                    
                    return(RC_INSERT_FAILED);
                }
                
                if(childsib_nleaf.write(sibPid, pf) != 0)
                {
                    // @ERROR: Could not write nonleaf node contents to 
                    //         pagefile
                    
                    return(RC_INSERT_FAILED);
                }
                
                if(nleaf.read(pid, pf) != 0)
                {
                    // @ERROR: Could not retrieve nonleaf node from pagefile
                    
                    return(RC_INSERT_FAILED);
                }
                
                if(nleaf.nleaf_pf.node.entryCount == 
                            BTNonLeafNode::MAX_NODE_ENTRIES)
                {
                    // The nonleaf node is full
                    
                    of_key = mid_key;
                    of_pid = sibPid;
                    
                    return(RC_NLEAF_OVERFLOW);
                }
                else
                {
                    // The key can be inserted into the nonleaf node
                    
                    if(nleaf.insert(mid_key, sibPid) != 0)
                    {
                        // @ERROR: Could not insert key into nonleaf node
                        
                        return(RC_INSERT_FAILED);
                    }
                }
                
                // The first key in the sibling nonleaf node was 
                // successfully inserted into the nonleaf node
                
                if(nleaf.write(pid, pf) != 0)
                {
                    // @ERROR: Could not write nonleaf node contents to 
                    //         pagefile
                    
                    return(RC_INSERT_FAILED);
                }
                
                return(0);
            }
            case(RC_INSERT_FAILED):
            default:
            {
                // @ERROR: A lower-level insert operation failed
                
                return(RC_INSERT_FAILED);
            }
        }
    }
}
